import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.linalg import solve_continuous_are

class WheeledLegRobot:
    def __init__(self):
        # Robot parameters
        self.m_body = 2.0      # body mass kg
        self.m_leg = 0.5       # leg mass kg
        self.l_hip = 0.3       # hip joint length m
        self.l_knee = 0.3      # knee joint length m
        self.I_body = 0.1      # body moment of inertia
        self.g = 9.81          # gravity
        
        # State: [x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot]
        self.state_dim = 12
        self.control_dim = 6
        
    def dynamics(self, state, control_torques, t, disturbance_type='none'):
        """Robot dynamics equations with disturbances"""
        x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot = state
        tau1, tau2, tau3, tau4, tau5, tau6 = control_torques
        
        # Simplified dynamics model
        M = np.diag([self.m_body, self.m_body, self.m_body, 
                     self.I_body, self.I_body, self.I_body])
        
        # Coriolis and gravity terms
        C = np.array([0, 0, self.m_body * self.g, 0, 0, 0])
        
        # Control input mapping
        B = np.eye(6)
        
        # Add disturbances based on type
        disturbance = self.get_disturbance(t, disturbance_type)
        
        # Acceleration calculation with disturbance
        accelerations = np.linalg.solve(M, B @ control_torques - C + disturbance)
        
        # State derivatives
        state_dot = np.zeros(12)
        state_dot[0:6] = state[6:12]  # position derivatives = velocities
        state_dot[6:12] = accelerations  # velocity derivatives = accelerations
        
        return state_dot
    
    def get_disturbance(self, t, disturbance_type):
        """Generate different types of disturbances"""
        disturbance = np.zeros(6)
        
        if disturbance_type == 'none':
            return disturbance
            
        elif disturbance_type == 'constant':
            # Constant disturbance force
            disturbance = np.array([2.0, 1.5, 3.0, 0.5, 0.5, 0.5])
            
        elif disturbance_type == 'sinusoidal':
            # Sinusoidal disturbance
            amplitude = 3.0
            frequency = 2.0
            disturbance = amplitude * np.array([
                np.sin(2*np.pi*frequency*t),
                np.cos(2*np.pi*frequency*t),
                0.5*np.sin(2*np.pi*frequency*t + np.pi/4),
                0.3*np.cos(2*np.pi*frequency*t),
                0.3*np.sin(2*np.pi*frequency*t),
                0.2*np.cos(2*np.pi*frequency*t)
            ])
            
        elif disturbance_type == 'impulse':
            # Impulse disturbances at specific times
            if 2.0 <= t < 2.1:
                disturbance = np.array([10, 8, 5, 2, 2, 2])
            elif 5.0 <= t < 5.1:
                disturbance = np.array([-8, 10, 3, 1, 1, 1])
            else:
                disturbance = np.zeros(6)
                
        elif disturbance_type == 'random':
            # Random noise disturbance
            disturbance = 2.0 * np.random.randn(6)
            
        elif disturbance_type == 'step':
            # Step disturbance
            if t >= 3.0:
                disturbance = np.array([4.0, 3.0, 2.0, 1.0, 1.0, 1.0])
                
        elif disturbance_type == 'complex':
            # Complex combined disturbance
            if t < 3.0:
                # Constant
                disturbance = np.array([1.0, 0.5, 2.0, 0.3, 0.3, 0.3])
            elif 3.0 <= t < 6.0:
                # Sinusoidal
                amplitude = 2.0
                disturbance = amplitude * np.array([
                    np.sin(2*np.pi*1.5*t),
                    np.cos(2*np.pi*1.5*t),
                    0,
                    0.2*np.sin(2*np.pi*1.5*t),
                    0.2*np.cos(2*np.pi*1.5*t),
                    0
                ])
            else:
                # Random
                disturbance = 1.5 * np.random.randn(6)
                
        return disturbance

class PIDController:
    def __init__(self, kp, ki, kd, dt):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.dt = dt
        self.integral = np.zeros(6)
        self.previous_error = np.zeros(6)
    
    def compute_control(self, desired_state, current_state):
        error = desired_state[0:6] - current_state[0:6]
        error_derivative = (error - self.previous_error) / self.dt
        
        self.integral += error * self.dt
        
        # Anti-windup
        self.integral = np.clip(self.integral, -50, 50)
        
        control = (self.kp * error + 
                  self.ki * self.integral + 
                  self.kd * error_derivative)
        
        self.previous_error = error.copy()
        return control

class SMCController:
    def __init__(self, lambda_val, k, phi, dt):
        self.lambda_val = lambda_val
        self.k = k
        self.phi = phi
        self.dt = dt
        self.previous_error = np.zeros(6)
    
    def compute_control(self, desired_state, current_state):
        error = desired_state[0:6] - current_state[0:6]
        error_derivative = (error - self.previous_error) / self.dt
        
        # Sliding surface
        s = error_derivative + self.lambda_val * error
        
        # Control law (using saturation function to reduce chattering)
        sat_s = np.tanh(s * 5)  # Reduced gain for better robustness
        control = self.k * sat_s + self.phi * s
        
        self.previous_error = error.copy()
        return control

class LQRController:
    def __init__(self, A, B, Q, R):
        # Use scipy to solve continuous-time algebraic Riccati equation
        P = solve_continuous_are(A, B, Q, R)
        self.K = np.linalg.solve(R, B.T @ P)
    
    def compute_control(self, desired_state, current_state):
        error = current_state - desired_state
        return -self.K @ error

def simulate_controller_with_disturbance(controller, robot, desired_trajectory, t_span, dt, disturbance_type):
    """Simulate specific controller with disturbance"""
    t_eval = np.arange(t_span[0], t_span[1], dt)
    n_steps = len(t_eval)
    
    # Initialize state records
    states = np.zeros((n_steps, robot.state_dim))
    controls = np.zeros((n_steps, robot.control_dim))
    errors = np.zeros((n_steps, 6))
    disturbances = np.zeros((n_steps, 6))
    
    # Initial state
    state = np.array([0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    
    for i, t in enumerate(t_eval):
        # Get desired state
        desired_state = desired_trajectory(t)
        
        # Calculate control input
        control = controller.compute_control(desired_state, state)
        
        # Limit control input
        control = np.clip(control, -25, 25)  # Increased limits for disturbance rejection
        
        # Get current disturbance
        current_disturbance = robot.get_disturbance(t, disturbance_type)
        
        # Record state and control
        states[i] = state
        controls[i] = control
        errors[i] = desired_state[0:6] - state[0:6]
        disturbances[i] = current_disturbance
        
        # Integrate dynamics with disturbance
        state_dot = robot.dynamics(state, control, t, disturbance_type)
        state = state + state_dot * dt
    
    return t_eval, states, controls, errors, disturbances

def desired_circular_trajectory(t):
    """Circular trajectory"""
    radius = 1.0
    omega = 0.5
    
    x = radius * np.cos(omega * t)
    y = radius * np.sin(omega * t)
    z = 0.5  # maintain fixed height
    theta, phi, psi = 0, 0, 0
    x_dot = -radius * omega * np.sin(omega * t)
    y_dot = radius * omega * np.cos(omega * t)
    z_dot = 0
    theta_dot, phi_dot, psi_dot = 0, 0, 0
    
    return np.array([x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot])

def desired_figure8_trajectory(t):
    """Figure-8 trajectory with varying height"""
    a = 1.0  # x-axis amplitude
    b = 0.7  # y-axis amplitude
    omega = 0.4  # frequency
    
    # Figure-8 pattern
    x = a * np.sin(omega * t)
    y = b * np.sin(2 * omega * t)
    
    # Varying height for more challenge
    z = 0.5 + 0.1 * np.sin(omega * t + np.pi/4)
    
    # Orientation - follow the trajectory direction
    theta = np.arctan2(y, x)  # body orientation follows movement direction
    phi = 0.1 * np.sin(omega * t)  # slight roll
    psi = 0.05 * np.cos(omega * t)  # slight yaw
    
    # Velocities
    x_dot = a * omega * np.cos(omega * t)
    y_dot = 2 * b * omega * np.cos(2 * omega * t)
    z_dot = 0.1 * omega * np.cos(omega * t + np.pi/4)
    
    # Angular velocities
    theta_dot = 0  # simplified
    phi_dot = 0.1 * omega * np.cos(omega * t)
    psi_dot = -0.05 * omega * np.sin(omega * t)
    
    return np.array([x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot])

def desired_spiral_trajectory(t):
    """3D spiral trajectory"""
    r0 = 0.5  # initial radius
    r_growth = 0.1  # radius growth rate
    omega = 0.6  # angular frequency
    z0 = 0.3  # initial height
    z_growth = 0.05  # height growth rate
    
    # Spiral coordinates
    r = r0 + r_growth * t
    x = r * np.cos(omega * t)
    y = r * np.sin(omega * t)
    z = z0 + z_growth * t
    
    # Orientation - tangent to spiral
    theta = np.arctan2(y, x)
    phi = 0.05 * np.sin(omega * t)
    psi = 0.02 * np.cos(omega * t)
    
    # Velocities
    x_dot = r_growth * np.cos(omega * t) - r * omega * np.sin(omega * t)
    y_dot = r_growth * np.sin(omega * t) + r * omega * np.cos(omega * t)
    z_dot = z_growth
    
    # Angular velocities
    theta_dot = 0  # simplified
    phi_dot = 0.05 * omega * np.cos(omega * t)
    psi_dot = -0.02 * omega * np.sin(omega * t)
    
    return np.array([x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot])

def desired_aggressive_trajectory(t):
    """Aggressive trajectory with sharp turns and varying speeds"""
    if t < 3:
        # First segment: fast linear motion
        x = 0.8 * t
        y = 0.2 * np.sin(2 * t)
        z = 0.4 + 0.1 * np.sin(t)
        x_dot = 0.8
        y_dot = 0.4 * np.cos(2 * t)
        z_dot = 0.1 * np.cos(t)
    elif t < 6:
        # Sharp turn
        x = 2.4 + 0.5 * np.cos((t-3) - np.pi)
        y = 0.2 + 0.5 * np.sin((t-3) - np.pi)
        z = 0.5 + 0.05 * np.sin(2*(t-3))
        x_dot = -0.5 * np.sin((t-3) - np.pi)
        y_dot = 0.5 * np.cos((t-3) - np.pi)
        z_dot = 0.1 * np.cos(2*(t-3))
    else:
        # Return with acceleration
        x = 1.9 - 0.3 * (t-6) + 0.1 * (t-6)**2
        y = -0.3 + 0.2 * (t-6)
        z = 0.45 + 0.08 * np.sin(3*(t-6))
        x_dot = -0.3 + 0.2 * (t-6)
        y_dot = 0.2
        z_dot = 0.24 * np.cos(3*(t-6))
    
    # Orientation follows movement direction
    theta = np.arctan2(y_dot, x_dot) if abs(x_dot) > 0.01 else 0
    phi = 0.08 * np.sin(1.5 * t)
    psi = 0.04 * np.cos(t)
    
    # Angular velocities
    theta_dot = 0  # simplified
    phi_dot = 0.12 * np.cos(1.5 * t)
    psi_dot = -0.04 * np.sin(t)
    
    return np.array([x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot])

def calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, disturbance_type):
    """Calculate robustness performance metrics"""
    
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    # Individual axis errors
    x_error_pid = np.abs(errors_pid[:, 0])
    y_error_pid = np.abs(errors_pid[:, 1])
    z_error_pid = np.abs(errors_pid[:, 2])
    
    x_error_smc = np.abs(errors_smc[:, 0])
    y_error_smc = np.abs(errors_smc[:, 1])
    z_error_smc = np.abs(errors_smc[:, 2])
    
    x_error_lqr = np.abs(errors_lqr[:, 0])
    y_error_lqr = np.abs(errors_lqr[:, 1])
    z_error_lqr = np.abs(errors_lqr[:, 2])
    
    # Steady-state error (last 2 seconds average)
    steady_start_idx = int(0.8 * len(error_norm_pid))
    
    # 1. 扰动抑制比 (Disturbance Rejection Ratio)
    def calculate_disturbance_rejection_ratio(error_norm, disturbance_type):
        """
        计算扰动抑制比
        在扰动期间的误差与无扰动期间误差的比值
        """
        if disturbance_type == 'none':
            return 1.0
            
        # 定义扰动期间
        disturbance_indices = []
        if disturbance_type == 'impulse':
            disturbance_indices = list(range(200, 210)) + list(range(500, 510))
        elif disturbance_type == 'step':
            disturbance_indices = list(range(300, len(error_norm)))
        elif disturbance_type == 'constant':
            disturbance_indices = list(range(0, len(error_norm)))
        else:
            # 对于其他扰动类型，使用误差较大的时期
            high_error_indices = np.where(error_norm > np.percentile(error_norm, 70))[0]
            disturbance_indices = high_error_indices
        
        if len(disturbance_indices) == 0:
            return 1.0
            
        # 无扰动期间的误差（前1秒）
        normal_error = np.mean(error_norm[:100])
        disturbance_error = np.mean(error_norm[disturbance_indices])
        
        if normal_error < 1e-6:
            return disturbance_error
            
        return disturbance_error / normal_error
    
    # 2. 误差积分指标 (Integral Error Metrics)
    def calculate_integral_metrics(error_norm):
        """计算各种误差积分指标"""
        # ITAE (时间乘绝对误差积分) - 惩罚持续误差
        itae = np.sum(np.arange(len(error_norm)) * 0.01 * error_norm)
        
        # ISE (误差平方积分) - 强调大误差
        ise = np.sum(error_norm**2)
        
        # IAE (绝对误差积分) - 总误差量
        iae = np.sum(error_norm)
        
        return itae, ise, iae
    
    # 3. 稳定时间 (Settling Time)
    def calculate_settling_time(error_norm, percentage=5):
        """计算稳定到最终值±percentage%的时间"""
        final_value = np.mean(error_norm[steady_start_idx:])
        if final_value < 1e-6:
            return 0.0
            
        threshold = final_value * (1 + percentage/100)
        
        # 找到最后一个超过阈值的时间点
        exceeding_indices = np.where(error_norm > threshold)[0]
        if len(exceeding_indices) == 0:
            return 0.0
            
        last_exceed_idx = exceeding_indices[-1]
        return last_exceed_idx * 0.01
    
    # 4. 鲁棒性裕度 (Robustness Margin)
    def calculate_robustness_margin(error_norm):
        """基于误差的波动程度计算鲁棒性裕度"""
        steady_state_error = error_norm[steady_start_idx:]
        if len(steady_state_error) == 0:
            return 0.0
            
        # 使用变异系数(标准差/均值)的倒数作为鲁棒性指标
        cv = np.std(steady_state_error) / (np.mean(steady_state_error) + 1e-6)
        return 1.0 / (1.0 + cv)  # 值越大表示鲁棒性越好
    
    # 计算各个指标
    disturbance_rejection_pid = calculate_disturbance_rejection_ratio(error_norm_pid, disturbance_type)
    disturbance_rejection_smc = calculate_disturbance_rejection_ratio(error_norm_smc, disturbance_type)
    disturbance_rejection_lqr = calculate_disturbance_rejection_ratio(error_norm_lqr, disturbance_type)
    
    itae_pid, ise_pid, iae_pid = calculate_integral_metrics(error_norm_pid)
    itae_smc, ise_smc, iae_smc = calculate_integral_metrics(error_norm_smc)
    itae_lqr, ise_lqr, iae_lqr = calculate_integral_metrics(error_norm_lqr)
    
    settling_time_pid = calculate_settling_time(error_norm_pid)
    settling_time_smc = calculate_settling_time(error_norm_smc)
    settling_time_lqr = calculate_settling_time(error_norm_lqr)
    
    robustness_margin_pid = calculate_robustness_margin(error_norm_pid)
    robustness_margin_smc = calculate_robustness_margin(error_norm_smc)
    robustness_margin_lqr = calculate_robustness_margin(error_norm_lqr)
    
    # 振荡指标
    oscillation_pid = np.std(error_norm_pid[steady_start_idx:])
    oscillation_smc = np.std(error_norm_smc[steady_start_idx:])
    oscillation_lqr = np.std(error_norm_lqr[steady_start_idx:])
    
    # 超调量
    steady_error_pid = np.mean(error_norm_pid[steady_start_idx:])
    steady_error_smc = np.mean(error_norm_smc[steady_start_idx:])
    steady_error_lqr = np.mean(error_norm_lqr[steady_start_idx:])
    
    overshoot_pid = (np.max(error_norm_pid) / steady_error_pid) if steady_error_pid > 0.001 else 0
    overshoot_smc = (np.max(error_norm_smc) / steady_error_smc) if steady_error_smc > 0.001 else 0
    overshoot_lqr = (np.max(error_norm_lqr) / steady_error_lqr) if steady_error_lqr > 0.001 else 0
    
    metrics = {
        'RMSE Position': [
            np.sqrt(np.mean(error_norm_pid**2)),
            np.sqrt(np.mean(error_norm_smc**2)),
            np.sqrt(np.mean(error_norm_lqr**2))
        ],
        'Max Position Error': [
            np.max(error_norm_pid),
            np.max(error_norm_smc),
            np.max(error_norm_lqr)
        ],
        'Steady-State Error': [
            steady_error_pid,
            steady_error_smc,
            steady_error_lqr
        ],
        'Disturbance Rejection Ratio': [
            disturbance_rejection_pid,
            disturbance_rejection_smc,
            disturbance_rejection_lqr
        ],
        'ITAE': [
            itae_pid,
            itae_smc,
            itae_lqr
        ],
        'ISE': [
            ise_pid,
            ise_smc,
            ise_lqr
        ],
        'Settling Time (s)': [
            settling_time_pid,
            settling_time_smc,
            settling_time_lqr
        ],
        'Robustness Margin': [
            robustness_margin_pid,
            robustness_margin_smc,
            robustness_margin_lqr
        ],
        'Oscillation Level': [
            oscillation_pid,
            oscillation_smc,
            oscillation_lqr
        ],
        'Overshoot Ratio': [
            overshoot_pid,
            overshoot_smc,
            overshoot_lqr
        ]
    }
    
    return metrics

def run_robustness_test(disturbance_type, trajectory_type='figure8'):
    """Run robustness test for all controllers with specified disturbance and trajectory"""
    print(f"\n{'='*60}")
    print(f"ROBUSTNESS TEST: {disturbance_type.upper()} DISTURBANCE")
    print(f"TRAJECTORY: {trajectory_type.upper()}")
    print(f"{'='*60}")
    
    # Initialize robot and controllers
    robot = WheeledLegRobot()
    dt = 0.01
    t_span = (0, 10)
    
    # Select trajectory
    if trajectory_type == 'figure8':
        desired_trajectory_func = desired_figure8_trajectory
    elif trajectory_type == 'spiral':
        desired_trajectory_func = desired_spiral_trajectory
    elif trajectory_type == 'aggressive':
        desired_trajectory_func = desired_aggressive_trajectory
    elif trajectory_type == 'circular':
        desired_trajectory_func = desired_circular_trajectory
    else:
        desired_trajectory_func = desired_figure8_trajectory  # default
    
    # Controllers with slightly tuned parameters for robustness
    pid = PIDController(
        kp=np.array([18, 18, 22, 12, 12, 12]),
        ki=np.array([0.2, 0.2, 0.8, 0.2, 0.2, 0.2]),
        kd=np.array([6, 6, 10, 4, 4, 4]),
        dt=dt
    )
    
    smc = SMCController(
        lambda_val=3.0,
        k=np.array([12, 12, 15, 8, 8, 8]),
        phi=np.array([3, 3, 4, 2, 2, 2]),
        dt=dt
    )
    
    # LQR Controller
    A = np.zeros((12, 12))
    A[0:6, 6:12] = np.eye(6)
    A[8, 8] = -0.2
    
    B = np.zeros((12, 6))
    B[6:12, 0:6] = np.eye(6) / robot.m_body
    B[9:12, 3:6] = np.eye(3) / robot.I_body
    
    Q = np.diag([15, 15, 25, 8, 8, 8, 2, 2, 3, 1, 1, 1])
    R = np.diag([0.05, 0.05, 0.05, 0.05, 0.05, 0.05])
    
    lqr = LQRController(A, B, Q, R)
    
    # Run simulations with disturbance
    print("Running simulations with disturbance...")
    t_pid, states_pid, controls_pid, errors_pid, dist_pid = simulate_controller_with_disturbance(
        pid, robot, desired_trajectory_func, t_span, dt, disturbance_type)
    
    t_smc, states_smc, controls_smc, errors_smc, dist_smc = simulate_controller_with_disturbance(
        smc, robot, desired_trajectory_func, t_span, dt, disturbance_type)
    
    t_lqr, states_lqr, controls_lqr, errors_lqr, dist_lqr = simulate_controller_with_disturbance(
        lqr, robot, desired_trajectory_func, t_span, dt, disturbance_type)
    
    # Plot results
    metrics = plot_robustness_comparison(t_pid, t_smc, t_lqr, 
                             states_pid, states_smc, states_lqr,
                             errors_pid, errors_smc, errors_lqr,
                             dist_pid, disturbance_type, trajectory_type)
    
    return metrics

def plot_robustness_comparison(t_pid, t_smc, t_lqr, states_pid, states_smc, states_lqr,
                             errors_pid, errors_smc, errors_lqr, disturbances, disturbance_type, trajectory_type='figure8'):
    """Plot robustness comparison results"""
    
    fig, axes = plt.subplots(3, 3, figsize=(18, 14))
    fig.suptitle(f'Robustness Test: {disturbance_type.upper()} Disturbance - {trajectory_type.upper()} Trajectory', 
                 fontsize=16, fontweight='bold')
    
    # Generate reference trajectory for plotting
    t_ref = np.linspace(0, 10, 500)
    if trajectory_type == 'figure8':
        x_ref = 1.0 * np.sin(0.4 * t_ref)
        y_ref = 0.7 * np.sin(0.8 * t_ref)
        z_ref = 0.5 + 0.1 * np.sin(0.4 * t_ref + np.pi/4)
    elif trajectory_type == 'spiral':
        r0, r_growth, omega = 0.5, 0.1, 0.6
        z0, z_growth = 0.3, 0.05
        r = r0 + r_growth * t_ref
        x_ref = r * np.cos(omega * t_ref)
        y_ref = r * np.sin(omega * t_ref)
        z_ref = z0 + z_growth * t_ref
    elif trajectory_type == 'aggressive':
        x_ref, y_ref, z_ref = [], [], []
        for t in t_ref:
            state = desired_aggressive_trajectory(t)
            x_ref.append(state[0])
            y_ref.append(state[1])
            z_ref.append(state[2])
        x_ref, y_ref, z_ref = np.array(x_ref), np.array(y_ref), np.array(z_ref)
    else:  # circular
        x_ref = np.cos(0.5 * t_ref)
        y_ref = np.sin(0.5 * t_ref)
        z_ref = 0.5 * np.ones_like(t_ref)
    
    # Position tracking with disturbance indication
    axes[0,0].plot(t_pid, states_pid[:, 0], 'r-', label='PID', linewidth=2)
    axes[0,0].plot(t_smc, states_smc[:, 0], 'g-', label='SMC', linewidth=2)
    axes[0,0].plot(t_lqr, states_lqr[:, 0], 'b-', label='LQR', linewidth=2)
    axes[0,0].plot(t_ref, x_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    
    # Shade disturbance periods
    if disturbance_type in ['impulse', 'step']:
        axes[0,0].axvspan(2.0, 2.1, alpha=0.2, color='red', label='Disturbance')
        if disturbance_type == 'impulse':
            axes[0,0].axvspan(5.0, 5.1, alpha=0.2, color='red')
        elif disturbance_type == 'step':
            axes[0,0].axvspan(3.0, 10, alpha=0.2, color='red')
    
    axes[0,0].set_ylabel('X Position (m)', fontsize=12, labelpad=10)
    axes[0,0].set_title('X-Axis Position Tracking', fontsize=12, pad=15)
    axes[0,0].legend(fontsize=10)
    axes[0,0].grid(True)
    axes[0,0].tick_params(axis='both', which='major', labelsize=10)
    
    # Y position
    axes[0,1].plot(t_pid, states_pid[:, 1], 'r-', label='PID', linewidth=2)
    axes[0,1].plot(t_smc, states_smc[:, 1], 'g-', label='SMC', linewidth=2)
    axes[0,1].plot(t_lqr, states_lqr[:, 1], 'b-', label='LQR', linewidth=2)
    axes[0,1].plot(t_ref, y_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    axes[0,1].set_ylabel('Y Position (m)', fontsize=12, labelpad=10)
    axes[0,1].set_title('Y-Axis Position Tracking', fontsize=12, pad=15)
    axes[0,1].legend(fontsize=10)
    axes[0,1].grid(True)
    axes[0,1].tick_params(axis='both', which='major', labelsize=10)
    
    # Z position
    axes[0,2].plot(t_pid, states_pid[:, 2], 'r-', label='PID', linewidth=2)
    axes[0,2].plot(t_smc, states_smc[:, 2], 'g-', label='SMC', linewidth=2)
    axes[0,2].plot(t_lqr, states_lqr[:, 2], 'b-', label='LQR', linewidth=2)
    axes[0,2].plot(t_ref, z_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    axes[0,2].set_ylabel('Z Position (m)', fontsize=12, labelpad=10)
    axes[0,2].set_title('Z-Axis Position Tracking', fontsize=12, pad=15)
    axes[0,2].legend(fontsize=10)
    axes[0,2].grid(True)
    axes[0,2].tick_params(axis='both', which='major', labelsize=10)
    
    # Position Error Norm
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    axes[1,0].plot(t_pid, error_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,0].plot(t_smc, error_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,0].plot(t_lqr, error_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,0].set_ylabel('Position Error Norm (m)', fontsize=12, labelpad=10)
    axes[1,0].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[1,0].set_title('Position Error Norm', fontsize=12, pad=15)
    axes[1,0].legend(fontsize=10)
    axes[1,0].grid(True)
    axes[1,0].tick_params(axis='both', which='major', labelsize=10)
    
    # Individual Axis Errors
    axes[1,1].plot(t_pid, np.abs(errors_pid[:, 0]), 'r-', label='PID X', linewidth=1.5, alpha=0.7)
    axes[1,1].plot(t_smc, np.abs(errors_smc[:, 0]), 'g-', label='SMC X', linewidth=1.5, alpha=0.7)
    axes[1,1].plot(t_lqr, np.abs(errors_lqr[:, 0]), 'b-', label='LQR X', linewidth=1.5, alpha=0.7)
    axes[1,1].plot(t_pid, np.abs(errors_pid[:, 1]), 'r--', label='PID Y', linewidth=1.5, alpha=0.7)
    axes[1,1].plot(t_smc, np.abs(errors_smc[:, 1]), 'g--', label='SMC Y', linewidth=1.5, alpha=0.7)
    axes[1,1].plot(t_lqr, np.abs(errors_lqr[:, 1]), 'b--', label='LQR Y', linewidth=1.5, alpha=0.7)
    axes[1,1].set_ylabel('X/Y Axis Errors (m)', fontsize=12, labelpad=10)
    axes[1,1].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[1,1].set_title('Individual Axis Errors', fontsize=12, pad=15)
    axes[1,1].legend(ncol=2, fontsize=9)
    axes[1,1].grid(True)
    axes[1,1].tick_params(axis='both', which='major', labelsize=10)
    
    # Z Axis Error
    axes[1,2].plot(t_pid, np.abs(errors_pid[:, 2]), 'r-', label='PID', linewidth=2)
    axes[1,2].plot(t_smc, np.abs(errors_smc[:, 2]), 'g-', label='SMC', linewidth=2)
    axes[1,2].plot(t_lqr, np.abs(errors_lqr[:, 2]), 'b-', label='LQR', linewidth=2)
    axes[1,2].set_ylabel('Z Axis Error (m)', fontsize=12, labelpad=10)
    axes[1,2].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[1,2].set_title('Z-Axis Height Error', fontsize=12, pad=15)
    axes[1,2].legend(fontsize=10)
    axes[1,2].grid(True)
    axes[1,2].tick_params(axis='both', which='major', labelsize=10)
    
    # XY trajectory
    axes[2,0].plot(states_pid[:, 0], states_pid[:, 1], 'r-', label='PID', linewidth=2, alpha=0.7)
    axes[2,0].plot(states_smc[:, 0], states_smc[:, 1], 'g-', label='SMC', linewidth=2, alpha=0.7)
    axes[2,0].plot(states_lqr[:, 0], states_lqr[:, 1], 'b-', label='LQR', linewidth=2, alpha=0.7)
    axes[2,0].plot(x_ref, y_ref, 'k--', label='Desired', linewidth=2)
    axes[2,0].set_xlabel('X (m)', fontsize=12, labelpad=10)
    axes[2,0].set_ylabel('Y (m)', fontsize=12, labelpad=10)
    axes[2,0].set_title(f'{trajectory_type.upper()} Trajectory Tracking', fontsize=12, pad=15)
    axes[2,0].legend(fontsize=10)
    axes[2,0].grid(True)
    axes[2,0].axis('equal')
    axes[2,0].tick_params(axis='both', which='major', labelsize=10)
    
    # Disturbance profile
    disturbance_norm = np.linalg.norm(disturbances, axis=1)
    axes[2,1].plot(t_pid, disturbance_norm, 'k-', linewidth=2)
    axes[2,1].set_ylabel('Disturbance Magnitude', fontsize=12, labelpad=10)
    axes[2,1].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[2,1].set_title('Applied Disturbance Profile', fontsize=12, pad=15)
    axes[2,1].grid(True)
    axes[2,1].tick_params(axis='both', which='major', labelsize=10)
    
    # 鲁棒性裕度比较
    metrics = calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, disturbance_type)
    robustness_margins = metrics['Robustness Margin']
    controllers = ['PID', 'SMC', 'LQR']
    
    bars = axes[2,2].bar(controllers, robustness_margins, color=['red', 'green', 'blue'], alpha=0.7)
    axes[2,2].set_ylabel('Robustness Margin', fontsize=12, labelpad=10)
    axes[2,2].set_xlabel('Controller Type', fontsize=12, labelpad=10)
    axes[2,2].set_title('Controller Robustness Margin', fontsize=12, pad=15)
    axes[2,2].grid(True, alpha=0.3)
    axes[2,2].set_ylim(0, 1)
    axes[2,2].tick_params(axis='both', which='major', labelsize=10)
    
    # 在柱状图上添加数值标签
    for bar, margin in zip(bars, robustness_margins):
        height = bar.get_height()
        axes[2,2].text(bar.get_x() + bar.get_width()/2, height + 0.02, 
                      f'{margin:.3f}', ha='center', va='bottom', fontsize=11, fontweight='bold')
    
    # 调整子图之间的间距
    plt.tight_layout(pad=3.0)
    plt.subplots_adjust(top=0.93)
    
    plt.show()
    
    return metrics

def main():
    """Main function to run all robustness tests"""
    
    # Test configurations
    disturbance_types = ['constant', 'sinusoidal', 'impulse', 'step', 'random', 'complex']
    trajectory_types = ['figure8', 'spiral', 'aggressive']  # 测试不同的轨迹
    
    all_results = {}
    
    for traj_type in trajectory_types:
        print(f"\n{'#'*70}")
        print(f"TESTING TRAJECTORY: {traj_type.upper()}")
        print(f"{'#'*70}")
        
        traj_results = {}
        for dist_type in disturbance_types:
            metrics = run_robustness_test(dist_type, traj_type)
            traj_results[dist_type] = metrics
            
            # Print results
            print(f"\nPerformance under {dist_type} disturbance:")
            print(f"{'Metric':<25} {'PID':<12} {'SMC':<12} {'LQR':<12}")
            print("-" * 65)
            for metric, values in metrics.items():
                if metric in ['RMSE Position', 'Max Position Error', 'Steady-State Error']:
                    print(f"{metric:<25} {values[0]:<12.4f} {values[1]:<12.4f} {values[2]:<12.4f}")
                elif metric in ['Disturbance Rejection Ratio', 'Robustness Margin']:
                    print(f"{metric:<25} {values[0]:<12.4f} {values[1]:<12.4f} {values[2]:<12.4f}")
                elif metric in ['ITAE', 'ISE']:
                    print(f"{metric:<25} {values[0]:<12.2f} {values[1]:<12.2f} {values[2]:<12.2f}")
                elif metric in ['Settling Time (s)']:
                    print(f"{metric:<25} {values[0]:<12.2f} {values[1]:<12.2f} {values[2]:<12.2f}")
        
        all_results[traj_type] = traj_results
    
    # 综合性能分析
    print_comprehensive_analysis(all_results)

def print_comprehensive_analysis(all_results):
    """Print comprehensive performance analysis across all trajectories"""
    print(f"\n{'='*80}")
    print("COMPREHENSIVE PERFORMANCE ANALYSIS ACROSS ALL TRAJECTORIES")
    print(f"{'='*80}")
    
    trajectory_types = list(all_results.keys())
    disturbance_types = list(all_results[trajectory_types[0]].keys())
    
    # 计算每个控制器在所有测试中的平均性能
    overall_metrics = {}
    
    for metric in ['RMSE Position', 'Robustness Margin', 'Disturbance Rejection Ratio', 'ITAE']:
        controller_scores = {0: [], 1: [], 2: []}  # PID, SMC, LQR
        
        for traj_type in trajectory_types:
            for dist_type in disturbance_types:
                metrics = all_results[traj_type][dist_type]
                for i in range(3):
                    controller_scores[i].append(metrics[metric][i])
        
        # 计算平均值
        overall_metrics[metric] = [
            np.mean(controller_scores[0]),
            np.mean(controller_scores[1]),
            np.mean(controller_scores[2])
        ]
    
    controllers = ['PID', 'SMC', 'LQR']
    
    print("\nOverall Performance Averages:")
    print(f"{'Metric':<25} {'PID':<12} {'SMC':<12} {'LQR':<12}")
    print("-" * 65)
    for metric, values in overall_metrics.items():
        if metric in ['RMSE Position', 'Robustness Margin', 'Disturbance Rejection Ratio']:
            print(f"{metric:<25} {values[0]:<12.4f} {values[1]:<12.4f} {values[2]:<12.4f}")
        else:
            print(f"{metric:<25} {values[0]:<12.2f} {values[1]:<12.2f} {values[2]:<12.2f}")
    
    # 最终排名
    print(f"\n{'='*50}")
    print("FINAL CONTROLLER RANKING")
    print(f"{'='*50}")
    
    # 基于多个指标的综合排名
    rmse_ranking = np.argsort(overall_metrics['RMSE Position'])
    robustness_ranking = np.argsort(overall_metrics['Robustness Margin'])[::-1]
    rejection_ranking = np.argsort(overall_metrics['Disturbance Rejection Ratio'])
    
    print("\nBased on Tracking Accuracy (RMSE):")
    for i, rank in enumerate(rmse_ranking):
        print(f"{i+1}. {controllers[rank]} (RMSE: {overall_metrics['RMSE Position'][rank]:.4f}m)")
    
    print("\nBased on Robustness Margin:")
    for i, rank in enumerate(robustness_ranking):
        print(f"{i+1}. {controllers[rank]} (Margin: {overall_metrics['Robustness Margin'][rank]:.4f})")
    
    print("\nBased on Disturbance Rejection:")
    for i, rank in enumerate(rejection_ranking):
        print(f"{i+1}. {controllers[rank]} (Ratio: {overall_metrics['Disturbance Rejection Ratio'][rank]:.4f})")

if __name__ == "__main__":
    main()